from pathlib import Path
import pandas as pd
import reader
import os

CONFIG_FILENAME = 'config.ini'
DEMAND_FILENAME = 'demand.csv'
QOS_FILENAME = 'qos.csv'
SITE_BANDWIDTH_FILE = 'site_bandwidth.csv'

OUTPUT_FILE_NAME = 'solution.txt'

def get_abs_path(a: str, b: str):
    return Path.resolve(Path(a, b))

def simulator(data_base_path: str, output_base_path: str, verbose=False):
    qos_cons = reader.get_config_from(get_abs_path(data_base_path, CONFIG_FILENAME))
    
    demands, qos, site_bandwidth = [reader.get_data_from(get_abs_path(data_base_path, filename)) for filename in (DEMAND_FILENAME, QOS_FILENAME, SITE_BANDWIDTH_FILE)]
    
    ans = reader.get_ans_from(get_abs_path(output_base_path, OUTPUT_FILE_NAME))
    
    demands = demands.drop(columns='mtime')
    qos = qos.set_index('site_name')
    site_bandwidth = site_bandwidth.set_index('site_name')
    bandwidth_seq = pd.DataFrame(columns=site_bandwidth.index)
    n_user = len(demands.columns)

    sites = list(site_bandwidth.index)
    for (index, demand), plan in zip(demands.iterrows(), list(zip(*[iter(ans)]*n_user))):
        plan = dict(plan)

        rec = {site: 0 for site in sites}
        for user, d in demand.iteritems():
            strategy = plan[user]
            
            assert (d == 0 and strategy is None) or (strategy is not None and d == sum(strategy.values())), f'user demand false, user {user} demand {d} at time {index} (start at 0), but got {sum(strategy.values()) if strategy is not None else None}'
            
            if strategy is None:
                continue
            
            for site in strategy.keys():
                assert qos.loc[site, user] < qos_cons, f'QOS constraint false, site {site} to user {user} at time {index} (start at 0)'
        
            for site, given in strategy.items():
                rec[site] += given
        
        for site, bandwidth in rec.items():
            assert bandwidth < site_bandwidth.loc[site, 'bandwidth'], f"bandwidth exceeded, site {site}, real bandwidth {bandwidth}, required bandwidth {site_bandwidth.loc[site, 'bandwidth']} at time {index} (start at 0)"
        
        bandwidth_seq = pd.concat([bandwidth_seq, pd.DataFrame([rec])], axis=0, ignore_index=True)
        
    # return sum(bandwidth_seq.quantile(0.95, numeric_only=False))
    
    if verbose:
        print(bandwidth_seq.quantile(0.95, numeric_only=False, interpolation='higher'))

        output_dir = '__temp'

        os.makedirs(output_dir, exist_ok=True)

        bandwidth_seq.quantile(0.95, numeric_only=False, interpolation='higher').to_csv(f'{output_dir}/bandwidth_quant.csv')

        print(f'detailed information is saved in {output_dir}/bandwidth_quant.csv')
    print('ans: ', sum(bandwidth_seq.quantile(0.95, numeric_only=False, interpolation='higher')))


if __name__=='__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='../data', help='data path')
    parser.add_argument('--ans', type=str, default='../output', help='output file path')
    parser.add_argument('--verbose', action='store_true')

    args = parser.parse_args()

    if args.verbose:
        print(args)
    
    # print(simulator(args.data, args.ans))
    simulator(args.data, args.ans, args.verbose)